import h5py
import numpy as np
from torch.utils.data import Dataset


class TrainDataset(Dataset):
    def __init__(self, h5_file):
        super(TrainDataset, self).__init__()
        self.h5_file = h5_file
        self.f = h5py.File(self.h5_file, 'r')
        self.lr = self.f['lr'][:]
        self.hr = self.f['hr'][:]

    def __getitem__(self, idx):
        return self.lr[idx], self.hr[idx]

    def __len__(self):
        return len(self.lr)


class EvalDataset(Dataset):
    def __init__(self, h5_file):
        super(EvalDataset, self).__init__()
        self.h5_file = h5_file
        self.f = h5py.File(self.h5_file, 'r')
        self.lr = np.array(self.f['lr'])
        self.hr = np.array(self.f['hr'])
        self.bicubic = np.array(self.f['bicubic'])

    def __getitem__(self, idx):
        return self.lr[idx], self.hr[idx], self.bicubic[idx]

    def __len__(self):
        return len(self.lr)
